# Written by Bram Cohen
# see LICENSE.txt for license information

from Encrypter import OPTION_PATTERN
from BitTornado.bitfield import Bitfield
from BitTornado.clock import clock
from BitTornado.bencode import bencode, bdecode, xbdecode
from binascii import b2a_hex
from BitTornado.__init__ import product_name, version_short
from random import shuffle
from socket import inet_aton, inet_ntoa
from types import StringType

DEBUG1 = False
DEBUG2 = False

def toint(s):
    return long(b2a_hex(s), 16)

def tobinary(i):
    return (chr(i >> 24) + chr((i >> 16) & 0xFF) + 
        chr((i >> 8) & 0xFF) + chr(i & 0xFF))

def tobinary16(i):
    return chr((i >> 8) & 0xFF) + chr(i & 0xFF)

# Message IDs
CHOKE = chr(0)
UNCHOKE = chr(1)
INTERESTED = chr(2)
NOT_INTERESTED = chr(3)
HAVE = chr(4)
BITFIELD = chr(5)
REQUEST = chr(6)
PIECE = chr(7)
CANCEL = chr(8)
PORT = chr(9)

EXTENDED = chr(20)
UT_PEX = 1
##UT_METADATA = 3

# Message length
MSG_LENGTH = { CHOKE: 1, UNCHOKE: 1, INTERESTED: 1, NOT_INTERESTED: 1,
               HAVE:5, REQUEST: 13, CANCEL: 13, PORT: 3,
               BITFIELD: None, PIECE: None, EXTENDED: None }


class DecodeError(Exception): pass

class Connection:
    def __init__(self, connection, connecter, ccount):
        self.connection = connection
        self.connecter = connecter
        self.ccount = ccount
        self.got_anything = False
        self.next_upload = None
        self.outqueue = []
        self.partial_message = ''
        self.download = None
        self.send_choke_queued = False
        self.just_unchoked = None
        self._hold = False

        self.lastActivity = clock()
        
        self.client_version = None
        self.prefer_encryption = False

        # Extended Extentions
        self.support_extended_ext = False
        self.got_extended_ext = False
        self.pex_extention = 0
##        self.metadata_extention = 0

    def get_ip(self, real=False):
        return self.connection.get_ip(real)

    def get_port(self):
        return self.connection.get_port()
    
    def set_port(self, port):
        self.connection.set_port(port)
        
    def get_id(self):
        return self.connection.get_id()

    def get_readable_id(self):
        return self.connection.get_readable_id()

    def close(self):
        if DEBUG1:
            print (self.ccount,'connection closed')
        self.connection.close()

    def is_locally_initiated(self):
        return self.connection.is_locally_initiated()

    def is_encrypted(self):
        return self.connection.is_encrypted()

    def does_prefer_encryption(self):
        return self.prefer_encryption or self.is_encrypted()
    
    # Protocol v1
    def send_interested(self):
        self._send_message(INTERESTED)

    def send_not_interested(self):
        self._send_message(NOT_INTERESTED)

    def send_choke(self):
        if self.partial_message:
            self.send_choke_queued = True
        else:
            self._send_message(CHOKE)
            self.upload.choke_sent()
            self.just_unchoked = 0

    def send_unchoke(self):
        if self.send_choke_queued:
            self.send_choke_queued = False
            if DEBUG1:
                print (self.ccount,'CHOKE SUPPRESSED')
        else:
            self._send_message(UNCHOKE)
            if ( self.partial_message or self.just_unchoked is None
                 or not self.upload.interested or self.download.active_requests ):
                self.just_unchoked = 0
            else:
                self.just_unchoked = clock()

    def send_request(self, index, begin, length):
        self._send_message(REQUEST + tobinary(index) + 
            tobinary(begin) + tobinary(length))
        if DEBUG1:
            print (self.ccount,'sent request',index,begin,begin+length)

    def send_cancel(self, index, begin, length):
        self._send_message(CANCEL + tobinary(index) + 
            tobinary(begin) + tobinary(length))
        if DEBUG1:
            print (self.ccount,'sent cancel',index,begin,begin+length)

    def send_bitfield(self, bitfield):
        if not bitfield.empty():
            self._send_message(BITFIELD + str(bitfield))

    def send_have(self, index):
        self._send_message(HAVE + tobinary(index))

    def send_keepalive(self):
        self._send_message('')

    def send_port(self):
        dht = self.connecter.bt1download.dht
        if dht is not None:
            self._send_message(PORT + tobinary16(dht.port))

    def get_upload(self):
        return self.upload

    def get_download(self):
        return self.download

    def set_download(self, download):
        self.download = download

    def backlogged(self):
        return not self.connection.is_flushed()

    def hold_messages(self):        # efficiency measure
        self._hold = True

    def release_messages(self):
        self._hold = False
        if not self.partial_message and self.outqueue:
            self.connection.send_message_raw(''.join(self.outqueue))
            self.outqueue = []

    def _send_message(self, s):
        if DEBUG2:
            if s:
                print (self.ccount,'SENDING MESSAGE',ord(s[0]),len(s))
            else:
                print (self.ccount,'SENDING MESSAGE',-1,0)
        s = tobinary(len(s))+s
        if self.partial_message or self._hold:
            self.outqueue.append(s)
        else:
            self.connection.send_message_raw(s)

    def send_partial(self, bytes):
        if self.connection.closed:
            return 0
        self.lastActivity = clock()
        if bytes <= len(self.partial_message):
            self.connection.send_message_raw(self.partial_message[:bytes])
            self.partial_message = self.partial_message[bytes:]
            return bytes
        if self.outqueue:
            q = [self.partial_message]
            q.extend(self.outqueue)
            self.outqueue = []
            self.partial_message = ''.join(q)
        while bytes > len(self.partial_message):
            s = self.upload.get_upload_chunk()
            if s is None:
                break
            index, begin, piece = s
            self.partial_message = ''.join((
                        self.partial_message,
                        tobinary(len(piece) + 9), PIECE,
                        tobinary(index), tobinary(begin), piece.tostring() ))
            if DEBUG1:
                print (self.ccount,'sending chunk',index,begin,begin+len(piece))
        if bytes <= len(self.partial_message):
            self.connection.send_message_raw(self.partial_message[:bytes])
            self.partial_message = self.partial_message[bytes:]
            return bytes
        if self.send_choke_queued:
            self.send_choke_queued = False
            self.partial_message += tobinary(1)+CHOKE
            self.upload.choke_sent()
            self.just_unchoked = 0
        if not self.partial_message:
            return 0
        if bytes < len(self.partial_message):
            self.connection.send_message_raw(self.partial_message[:bytes])
            self.partial_message = self.partial_message[bytes:]
            return bytes
        self.connection.send_message_raw(self.partial_message)
        l = len(self.partial_message)        
        self.partial_message = ''
        return l


    def got_request(self, i, p, l):
        self.upload.got_request(i, p, l)
        if self.just_unchoked:
            self.connecter.ratelimiter.ping(clock() - self.just_unchoked)
            self.just_unchoked = 0

    # EXTENDED EXTENTION
    def send_extended(self):
        """
        extended handshake message
        """
        assert self.support_extended_ext

        extended_msg = {'m': {},
                        'v': product_name + " " + version_short,
                        'e': self.connecter.config['crypto_allowed']}
        
        if not self.connecter.bt1download.private:
            # ut pex
            if self.connecter.config['pex_allowed']:
                extended_msg['m']["ut_pex"] = UT_PEX
##            # ut metadata exchange
##            extended_msg['m']["ut_metadata"] = UT_METADATA
##            if self.connecter.metadataHandler.HasMetadata():
##                extended_msg['metadata_size'] = self.connecter.metadataHandler.getLength()

        if self.connecter.bt1download.port is not None:
            extended_msg['p'] = self.connecter.bt1download.port
        s = EXTENDED + chr(0) + bencode(extended_msg)
        self._send_message(s)
        
    def got_extended(self, info):
        if type(info) is not dict:
            return
        self.got_extended_ext = True
        # Additional messages
        if info.has_key('m'):
            # Enable/Disable ut PEX & metadata exchange
            if info['m'].has_key('ut_pex'):
                self.set_extention("pex_extention", info['m']['ut_pex'])
##            if info['m'].has_key('ut_metadata'):
##                self.set_extention("metadata_extention", info['m']['ut_metadata'])

        # Version name
        if info.has_key('v') and info['v'] != ('\xb5Torrent 1.5'):
            self.client_version = info['v']
        # Port Number
        if info.get('p'):
            self.set_port(info['p'])
        # Enable/Disable prefer encryption
        if info.get('e'):
            self.prefer_encryption = True
        else:
            self.prefer_encryption = False
##        # ut metadata exchange
##        if info.get('metadata_size') and not self.connecter.metadataHandler.HasMetadata():
##            self.connecter.metadataHandler.on_metadata_size(info['metadata_size'])

    def set_extention(self, extention, msgId):
        if type(msgId) != StringType:
            msgId = chr(msgId)
        setattr(self, extention, msgId)
            
    def send_ut_pex_msg(self, pexMessage):
        """
        Send utorrent style pex message
        """
        if not self.pex_extention:
            print "Error: sending pex while not allowed"
            return

        # Check the message
        if not pexMessage:
            return

        # Send the Message
        s = EXTENDED + self.pex_extention + bencode(pexMessage)
        self._send_message(s)

        if DEBUG1:
            print (self.ccount, "PEX Message Sent")    
        
    def got_ut_pex_msg(self, info):
        if type(info) is not dict or \
           not info.get('added'):
            return
        
        added = self._cut_already_known_peers(self.decode_peers(info['added'], info.get('added.f')))
        if added:
            # Maximum of 20 peers
            self.connection.Encoder.start_connections(added[:20])
        if DEBUG1:
            print (self.ccount, "Got PEX Message. peers: ", len(added))
            
    def decode_peers(self, added, flags):
        if len(added) % 6:
            return []

        peers = []
        x = 0
        f = 0
        while len(added) > x:
            encrypted = None
            seeder = False
            if flags:
                if ord(flags[f]) & 0x01:
                    encrypted = True
                if ord(flags[f]) & 0x02:
                    seeder = True
                f += 1
            if not (self.download.have.complete() and seeder):
                peers.append(((inet_ntoa(added[x:x+4]), int(toint(added[x+4:x+6]))), 0, encrypted))
            x += 6
        return peers
    
    def _cut_already_known_peers(self, list):
        connections = {}
        for c in self.connection.Encoder.connections.values(): 
            connections[c.get_ip()] = c
        for peer in list:
            peer_ip, peer_port = peer[0]
            if peer_ip in connections.keys():
                list.remove(peer)
                try:
                    if not c.get_port():
                        c.set_port(peer_port)            
                except:
                    pass

        return list

##    def send_ut_metadata_msg(self, mType, piece):
##        if not self.metadata_extention:
##            print "Error: sending ut_metadata while not allowed"
##            return
##        # Get the message 
##        msg = self.connecter.metadataHandler.write_metadata_packet(mType, piece)
##
##        # Send the Message
##        s = EXTENDED + self.metadata_extention + msg
##        print "send_ut_metadata_msg", s
##        self._send_message(s)
##        
##    def got_ut_metadata_msg(self, info, payload):
##        print "got_ut_metadata_msg", info
##        if type(info) is not dict \
##           or not info.has_key("msg_type") \
##           or not info.has_key("piece"):
##            print "Invalid ut_metadata_msg", info
##            return
##
##        mType = info["msg_type"]        
##        piece = info["piece"]
##        if mType == 0:      # request
##            if self.connecter.metadataHandler.HasMetadata():
##                self.send_ut_metadata_msg(1, piece)
##            else:
##                self.send_ut_metadata_msg(2, piece)
##        elif mType == 1:    # metadata
##            data = payload
##            total_size = info["total_size"]
##            self.connecter.metadataHandler.on_received_metadata(data, len(data), piece, total_size)
##        elif mType == 2:    # have no data
##            self.connecter.metadataHandler.on_recieved_nodata(piece)
##        else:
##            print "Invalid ut_metadata_msg type: %d" % mType
                    
class Connecter:
    def __init__(self, bt1download, choker, totalup, config, ratelimiter,
                 make_upload, downloader, numpieces, sched):
        self.bt1download = bt1download
#        self.metadataHandler = metadataHandler
        self.choker = choker
        self.totalup = totalup
        self.config = config
        self.ratelimiter = ratelimiter
        self.sched = sched
        self.rate_capped = False
        self.downloader = downloader
        self.make_upload = make_upload
        self.numpieces = numpieces
        
        self.connections = {}
        self.external_connection_made = 0
        self.ccount = 0
        
        if not self.bt1download.private and OPTION_PATTERN[43]:
            # pex messages every minute
            if self.config['pex_allowed']:
                self.ut_pex_sent = []
                self.sched(self.send_ut_pex, 60)
##            # metadata requests
##            if not self.metadataHandler.HasMetadata():
##                self.sched(self.send_ut_metadata, 10)

        #optimistic disconnect
        self.optimistic_disconnect_timeout = self.config['optimistic_disconnect_timeout']
        if self.optimistic_disconnect_timeout:
            self.optimistic_disconnect_timeout = max(300, self.optimistic_disconnect_timeout)
            self.sched(self.optimistic_disconnect , 60)

    def optimistic_disconnect(self):
        self.sched(self.optimistic_disconnect , 60)
        if self.downloader.paused:
            return
        if not self.bt1download.encoder.is_maximum_reached():
            return
        connections = list(self.connections.items())
        connections.sort(key = lambda x: x[1].lastActivity)
        now = clock()
        for c, con in connections:
            if now - con.lastActivity < self.optimistic_disconnect_timeout:
                break
            print "closing:", con.get_ip(), now - con.lastActivity
            c.close()
        
    # Extensions
    def send_ut_pex(self):
        self.sched(self.send_ut_pex, 60)
        if self.downloader.paused:
            return
        
        peers = []
        pex_peers = []
        for c in self.connections.values():
            if c.pex_extention:
                pex_peers.append(c)
            if c.get_port():
                peers.append(c)
                
        pexMessage = self.create_ut_pex_msg(peers)
        if pexMessage:
            for c in pex_peers:
                c.send_ut_pex_msg(pexMessage)
        self.ut_pex_sent = peers
            
    def create_ut_pex_msg(self, peers, max = 15):
        """
        Creates utorrent style pex message
        added   :   peers I connected to since last pex 
        added.f :   Flags for each peer in added (1=encrypted,2=seeder)
        dropped :   peers I disconnected from since last pex
        """
        def _encode(peers):
            msg = ""
            for peer in peers:
                msg += inet_aton(peer.get_ip()) + tobinary16(peer.get_port())
            return msg
        def _flags_encode(peers):
            flags = ""
            for peer in peers:
                sign = 0x00
                if peer.does_prefer_encryption():
                    sign |= 0x01
                if peer.download.have.complete():
                    sign |= 0x02
                flags += chr(sign)
            return _encode(peers), flags
        
        added = [peer for peer in peers if (peer not in self.ut_pex_sent)]
        dropped = [peer for peer in self.ut_pex_sent if (peer not in peers)]
        
        if len(added) > max:
            shuffle(added)
            added = added[:max]
            
        msg = {}
        if added:
            msg["added"], msg["added.f"] = _flags_encode(added)
##        if dropped:
##            msg["dropped"] = _encode(dropped)
        return msg

##    def send_ut_metadata(self):
##        if self.downloader.paused:
##            self.sched(self.send_ut_metadata, 10)
##            return
##        for c in self.connections.values():
##            if c.metadata_extention:
##                piece =  self.metadataHandler.next()
##                if piece is None:
##                    break
##                c.send_ut_metadata_msg(0, piece)
##                self.sched(lambda:self.metadataHandler.on_recieved_nodata_timeout(piece), 15)
##                
##        if not self.metadataHandler.HasMetadata(): 
##            self.sched(self.send_ut_metadata, 30)

    # Connections
    def reinitialize_connections(self):        
        for c in  self.connections.values():
            b = c.download.have
            b.resize(self.numpieces)
            c.download.got_have_bitfield(b)
            
    def how_many_connections(self):
        return len(self.connections)

    def connection_made(self, connection):
        self.ccount += 1
        c = Connection(connection, self, self.ccount)
        if DEBUG2:
            print (c.ccount,'connection made')
            
        self.connections[connection] = c
        
        # Send the Extended handshake immediately 
        if connection.options[43] and OPTION_PATTERN[43]:
            c.support_extended_ext = True
            c.send_extended()
 		
        # Make upload/download
        c.upload = self.make_upload(c, self.ratelimiter, self.totalup)
        c.download = self.downloader.make_download(c)

        # tell choker
        self.choker.connection_made(c)

        # Send the PORT
        if connection.options[63] and OPTION_PATTERN[63]:
            c.send_port()

        return c

    def connection_lost(self, connection):
        c = self.connections[connection]
        if DEBUG2:
            print (c.ccount,'connection closed')
        del self.connections[connection]
        if c.download:
            c.download.disconnected()
        self.choker.connection_lost(c)

    def connection_flushed(self, connection):
        conn = self.connections[connection]
        if conn.next_upload is None and (conn.partial_message is not None
               or len(conn.upload.buffer) > 0):
            self.ratelimiter.queue(conn)
            
    def shutdown(self):
        self.connections.clear()

    # Pieces
    def got_piece(self, i):
        for co in self.connections.values():
            co.send_have(i)
            
    def _decode_piece(self, s):
        i = toint(s)
        if i >= self.numpieces:
            raise DecodeError
        return i

    # Incoming messages
    def _got_choke(self, c, message):
        c.download.got_choke()
        
    def _got_unchoke(self, c, message):
        c.download.got_unchoke()
        
    def _got_interested(self, c, message):
        if not c.download.have.complete():
            c.upload.got_interested()
            
    def _got_notinterested(self, c, message):
        c.upload.got_not_interested()
        
    def _got_have(self, c, message):
        if c.download.got_have(self._decode_piece(message[1:])):
            c.upload.got_not_interested()

##    def _got_have_nometadata(self, c, message):
##        i = toint(message[1:])
##        if i > 65535: #index > 64k, just ignore it
##            return 
##        if len(c.download.have) < i:
##            c.download.have.resize(i+1)
##        try:
##            if c.download.got_have(i):
##                c.upload.got_not_interested()
##        except:
##            print i, len(c.download.have)
            
    def _got_bitfield(self, c, message):
        try:
            b = Bitfield(self.numpieces, message[1:])
        except ValueError:
            if DEBUG2:
                print (c.ccount,'bad bitfield')
            return False
        if c.download.got_have_bitfield(b):
            c.upload.got_not_interested()
            
##    def _got_bitfield_nometadata(self, c, message):
##        # We can't verify bitfield length
##        # for now save spare bits as well
##        length = len(message[1:]) * 8
##        try:
##            b = Bitfield(length, message[1:])
##        except ValueError:
##            print (c.ccount,'metadataless: bad bitfield')
##            return False
##        if c.download.got_have_bitfield(b):
##            c.upload.got_not_interested()
            
    def _got_request(self, c, message):
        c.got_request( self._decode_piece(message[1:5]),
                       toint(message[5:9]),toint(message[9:]) )
        
    def _got_cancel(self, c, message):
        c.upload.got_cancel( self._decode_piece(message[1:5]),
                             toint(message[5:9]),toint(message[9:]) )
        
    def _got_piece(self, c, message):
        if len(message) <= 9:
            if DEBUG2:
                print (c.ccount,'bad message length')
            return False
        c.lastActivity = clock()
        i = self._decode_piece(message[1:5])
        if c.download.got_piece(i, toint(message[5:9]), message[9:]):
            self.got_piece(i)
                                
    def _got_port(self, c, message):
        dht = self.bt1download.dht
        if dht is not None:
            dht.addContact(c.get_ip(), toint(message[1:3]))

    def _got_extended(self, c, message):
##        if message[1] == chr(UT_METADATA):
##            print "UT_METADATA", [message]

        try:
            info, length = xbdecode(message[2:], sloppy = 1)
        except ValueError:
            if DEBUG2:
                print (c.ccount,'bad EXTENDED message')
            return False            
        if message[1] == chr(0):
            c.got_extended(info)                        
        elif message[1] == chr(UT_PEX):
            c.got_ut_pex_msg(info)
##        elif message[1] == chr(UT_METADATA):
##            c.got_ut_metadata_msg(info, message[length+2:])
        elif DEBUG2:
            print (c.ccount, 'unknown extension message', ord(message[1]))

    # Map incoming messages
    def got_message(self, connection, message):
        c = self.connections[connection]
##        if self.metadataHandler.HasMetadata():
##            r = self._got_message(c, message, self._funcmap)
##        else:
##            r = self._got_message(c, message, self._funcmap_nometadata)
        r = self._got_message(c, message, self._funcmap)
        if not r:
            connection.close()

##    _funcmap_nometadata = { CHOKE: _got_choke,
##                            UNCHOKE: _got_unchoke,
##                            INTERESTED: _got_interested,
##                            NOT_INTERESTED: _got_notinterested,
##                            HAVE: _got_have_nometadata,
##                            BITFIELD: _got_bitfield_nometadata,
##                            #REQUEST: _got_request,
##                            #CANCEL: _got_cancel,
##                            #PIECE: _got_piece,
##                            PORT: _got_port,
##                            EXTENDED: _got_extended,
##                            }        
    _funcmap = { CHOKE: _got_choke,
                 UNCHOKE: _got_unchoke,
                 INTERESTED: _got_interested,
                 NOT_INTERESTED: _got_notinterested,
                 HAVE: _got_have,
                 BITFIELD: _got_bitfield,
                 REQUEST: _got_request,
                 CANCEL: _got_cancel,
                 PIECE: _got_piece,
                 PORT: _got_port,
                 EXTENDED: _got_extended,
                 }
    def _got_message(self, c, message, funcmap):
        t = message[0]
        if DEBUG2:
            print (c.ccount,'message received',ord(t))
        try:
            f = funcmap[t]
        except (KeyError, ValueError):
            if DEBUG2:
                print (c.ccount,'unknown message')
            return False
        l = MSG_LENGTH[t]
        if l and l != len(message):
            if DEBUG2:
                print (c.ccount,'bad message length')
            return False
        if c.got_anything:
            if t == BITFIELD:
                if DEBUG2:
                    print (c.ccount,'misplaced bitfield')
                return False
        elif t != EXTENDED and t != PORT:
            c.got_anything = True
            if t != BITFIELD:
                c.download.got_have_bitfield(Bitfield(self.numpieces)) # empty
        try:
            return f(self,c,message) is None
        except DecodeError:
            if DEBUG2:
                print (c.ccount,'bad piece number')
            return False
